"""

做一个基于函数的深度学习识别的杀毒软件

第一步：
提取出函数

第二步：
把提取出来的函数做处理
处理思路：
遇到的问题主要是不同程序的函数的对应关系该怎么去做？
做法：列表匹配


将常见/就按照最开始的这个程序的函数提取出来作为序列 完成

单个数据集处理
接下来就是多个数据集的处理


数据集大概结构：
横向是函数，纵向是标签
[[[1,1,1,1,0],[0]],[[0,1,1,1,0],[1]]]
用序列还是上面这个
['78 20 26','1 2 41']
用下面这个

下一步工作：
把上面这个np array弄出来
选择算法和结构
考虑输入输出

np array  完成
考虑输入输出 完成


接下来就是算法和结构

"""

import pefile
import sys
import random
import os

def chakanhanshu(path):
    pe = pefile.PE(path)

    print(pe.DIRECTORY_ENTRY_IMPORT[0].imports[0].name)
    #f = open("apis.txt",'a')
    #f.write("1"+'\r\n')
    #f.close()

    for inputdlls in pe.DIRECTORY_ENTRY_IMPORT:
        #print(inputdlls.dll)
        for api in inputdlls.imports:
            print(api.name)
            f = open(path+".txt",'a')
            f.write(api.name.decode("utf-8")+'\r')
            f.close()
    return path

def chuanjianchushixiangliang(path):
    liebiao = []
    with open(path+'.txt') as fp:
        for i, line in enumerate(fp):
            newline = line.replace('\n','')
            liebiao.append(newline)
    print(liebiao)
    return liebiao


def duibihanshu(liebiao,liebiao2):
    list1 = liebiao
    list2 = liebiao2
    list3 = set(list1)&set(list2)
    f = open("normal.txt",'w')
    list4 = [0]*79
    list5 = ''
    for i in list3:
        f.write(str(liebiao.index(i))+' ')
        list5 += str(liebiao.index(i))+" "
        suoying = liebiao.index(i)
        list4[suoying] = 1
    return list4,list5





# 删除空文件夹，否则会报错
def traverse(filepath):
    # 遍历filepath下所有文件，包括子目录
    files = os.listdir(filepath)
    for fi in files:
        fi_d = os.path.join(filepath, fi)
        if os.path.isdir(fi_d):  # 判断是否为文件夹
            if not os.listdir(fi_d):  # 如果文件夹为空
                os.rmdir(fi_d)  # 删除这个空文件夹
            else:
                traverse(fi_d)
        else:
            file = os.path.join(filepath, fi_d)
            #if os.path.getsize(file) == 0:  # 文件大小为0
                #os.remove(file)  # 删除这个文件


def get_file(path):  # 创建一个空列表
    files = os.listdir(path)
    files.sort()  # 排序
    list = []
    for file in files:
        if not os.path.isdir(path + file):  # 判断该文件是否是一个文件夹
            f_name = str(file)
            #             print(f_name)
            tr = '\\'  # 多增加一个斜杠
            filename = path + tr + f_name
            list.append(filename)
    return (list)

def main(filename,shuxing):
    yangben = filename
    shuxing = shuxing
    chakanhanshu(r'class4.exe')
    liebiao = chuanjianchushixiangliang(r'class4.exe')
    chakanhanshu(yangben)
    liebiao2 = chuanjianchushixiangliang(yangben)
    shujuji1,shujuji2 = duibihanshu(liebiao,liebiao2)
    print(shujuji1)
    print('\r')
    print(shujuji2)
    if shuxing == "virus":
        suijishu1 = random.randint(0,9999999999999999999)
        f = open('virus/'+str(suijishu1)+'.txt','a')
        f.write(shujuji2)
        f.close()
    if shuxing == "normal":
        suijishu2 = random.randint(0, 9999999999999999999)
        f = open('normal/'+str(suijishu2)+'.txt','a')
        f.write(shujuji2)
        f.close()
    os.system('del class4.exe.txt')
    os.system('del '+yangben+'.txt')
    os.system('del normal.txt')



if __name__ == '__main__':
    path = r'virus_file'
    traverse(path)
    list = get_file(path)
    print(list)
    for filename in list:
        main(filename,"virus")
    path = r'normal_file'
    traverse(path)
    list = get_file(path)
    for filename in list:
        main(filename,"normal")


